# -*- coding: utf-8 -*-
# @Author  : gaoyu
# @Time    : 2024/2/7 9:36 AM
# @Function:
class Muskingum:
    """
    呈村流域河道汇流
    """

    def __init__(self):
        self.I2 = None
        self.I1 = None
        self.dt = None
        self.O = []
        self.O2 = None
        self.XE = None
        self.KE = None

    def SetParmameter(self, parameter):
        """
        设置参数
        :param parameter:
        :return:
        """
        self.KE = parameter.m_KE
        self.XE = parameter.m_XE

    def SetState(self, state):
        self.I1 = state.m_QU0
        self.I2 = state.m_QU
        self.O = state.m_O
        self.dt = state.m_dt

    def UpdateState(self, state):
        state.m_O = self.O
        state.m_O2 = self.O2

    def calculate(self):
        KL = float(self.dt)  # 为了保证马斯京根法的两个线性条件，每个单元河取 KL = ∆t
        N = int(float(self.KE) / float(KL))  # 单元河段数
        XL = 0.5 - N * (1 - 2 * float(self.XE)) / 2  # 计算单元河段XL
        denominator = 0.5 * \
                      float(self.dt) + \
                      KL - \
                      KL * XL  # 计算分母
        C0 = (0.5 * float(self.dt) - KL * XL) / denominator
        C1 = (0.5 * float(self.dt) + KL * XL) / denominator
        C2 = (-0.5 * float(self.dt) + KL - KL * XL) / denominator

        if not self.O:
            self.O = list(range(N))  # 创建存储单元流域在子河段出口断面的出流量的动态数组

            for n in range(N):
                self.O[n] = 0.0  # 单元流域在子河段出口断面的出流量为0

        for n in range(N):
            O1 = self.O[n]  # 子河段时段初出流量
            O2 = C0 * self.I2 + C1 * self.I1 + C2 * O1  # 计算时段末单元流域在子河段出口断面的出流量，m3/s
            self.O[n] = O2  # 更新子河段时段初出流量
            self.I1 = O1  # 上一河段时段初出流为下一河段时段初入流
            self.I2 = O2  # 上一河段时段末出流为下一河段时段末入流
